import numpy as np
import random
import matplotlib.pyplot as plt

def initialize_centroids(data, k):
    """
    初始化质心
    """
    return [data[random.randint(0, len(data) - 1)] for _ in range(k)]

def assign_clusters(data, centroids):
    """
    将数据点分配到最近的质心
    """
    clusters = {i: [] for i in range(k)}
    for features in data:
        distances = [np.linalg.norm(features - centroid) for centroid in centroids]
        cluster_id = distances.index(min(distances))
        clusters[cluster_id].append(features)
    return clusters

def calculate_new_centroids(clusters):
    """
    重新计算质心
    """
    return [np.mean(cluster, axis=0) for cluster in clusters.values()]

def k_means(data, k, max_iter=100):
    """
    k-means算法主函数
    """
    centroids = initialize_centroids(data, k)
    for _ in range(max_iter):
        clusters = assign_clusters(data, centroids)
        new_centroids = calculate_new_centroids(clusters)
        if np.array_equal(centroids, new_centroids):
            break
        centroids = new_centroids
    return clusters, centroids

# 加载Wine数据集
def load_wine_data(filename):
    """
    加载Wine数据集
    """
    with open(filename, 'r') as file:
        data = np.loadtxt(file, delimiter=',', skiprows=1)
    return data[:, 1:]  # 排除第一列的类别标签

# Wine数据集文件路径
wine_data_file = 'wine.data'

# 加载数据
wine_data = load_wine_data(wine_data_file)

# 选择k值
k = 5

# 运行k-means算法
clusters, centroids = k_means(wine_data, k)
# 可视化聚类结果
def plot_clusters(clusters, centroids, data):
    plt.figure(figsize=(12, 8))
    colors = ['r', 'g', 'b', 'y', 'c', 'm', 'orange', 'purple', 'brown']
    markers = ['o', 's', '^', 'p', '*', '+', 'x', 'd', 'v']
    
    for i, (cluster_id, cluster) in enumerate(clusters.items()):
        for features in cluster:
            plt.scatter(features[0], features[1], c=colors[i], marker=markers[i], label=f'Cluster {i}')
        plt.scatter(centroids[i][0], centroids[i][1], c='black', marker='x', s=100, label=f'Centroid {i}')
    
    plt.title('K-Means Clustering')
    plt.xlabel('Feature 1')
    plt.ylabel('Feature 2')
    plt.legend()
    plt.show()

# 可视化聚类结果
plot_clusters(clusters, centroids, wine_data)
print("Clusters:", clusters)
print("Centroids:", centroids)